package net.cyue.ort.llm.sampling;

import net.cyue.ort.llm.generator.GenerationConfig;

/**
 * 采样策略配置类
 * 包含采样相关的所有配置参数，与生成配置解耦
 */
public class SamplingConfig {
    private final float temperature;
    private final float topP;
    private final int topK;
    private final float repetitionPenalty;
    private final boolean doSample;
    private final long[] eosTokenIDs;
    private final int contextWindow;
    private final int minSequenceLength;
    private final int targetSequenceLength;
    
    public SamplingConfig(
        float temperature,
        float topP,
        int topK,
        float repetitionPenalty,
        boolean doSample,
        long[] eosTokenIDs,
        int contextWindow,
        int minSequenceLength,
        int targetSequenceLength
    ) {
        this.temperature = temperature;
        this.topP = topP;
        this.topK = topK;
        this.repetitionPenalty = repetitionPenalty;
        this.doSample = doSample;
        this.eosTokenIDs = eosTokenIDs;
        this.contextWindow = contextWindow;
        this.minSequenceLength = minSequenceLength;
        this.targetSequenceLength = targetSequenceLength;
    }
    
    /**
     * 默认配置
     */
    public static SamplingConfig defaultConfig() {
        return new SamplingConfig(
            0.8f,      // temperature
            0.8f,      // topP
            20,        // topK
            1.1f,      // repetitionPenalty
            true,      // doSample
            null,      // eosTokenIDs
            384,       // contextWindow
            20,        // minSequenceLength
            128        // targetSequenceLength
        );
    }
    
    /**
     * 从LLMGenerationConfig创建SamplingConfig
     */
    public static SamplingConfig fromGenerationConfig(GenerationConfig genConfig) {
        return new SamplingConfig(
            genConfig.getTemperature(),
            genConfig.getTopP(),
            genConfig.getTopK(),
            genConfig.getRepetitionPenalty(),
            genConfig.isDoSample(),
            genConfig.getEosTokenIDs(),
            384,       // contextWindow (默认值)
            20,        // minSequenceLength (默认值)
            128        // targetSequenceLength (默认值)
        );
    }
    
    /**
     * 构建器模式
     */
    public static class Builder {
        private float temperature = 0.8f;
        private float topP = 0.8f;
        private int topK = 20;
        private float repetitionPenalty = 1.1f;
        private boolean doSample = true;
        private long[] eosTokenIDs = null;
        private int contextWindow = 384;
        private int minSequenceLength = 20;
        private int targetSequenceLength = 128;
        
        public Builder temperature(float temperature) {
            this.temperature = temperature;
            return this;
        }
        
        public Builder topP(float topP) {
            this.topP = topP;
            return this;
        }
        
        public Builder topK(int topK) {
            this.topK = topK;
            return this;
        }
        
        public Builder repetitionPenalty(float repetitionPenalty) {
            this.repetitionPenalty = repetitionPenalty;
            return this;
        }
        
        public Builder doSample(boolean doSample) {
            this.doSample = doSample;
            return this;
        }
        
        public Builder eosTokenIDs(long[] eosTokenIDs) {
            this.eosTokenIDs = eosTokenIDs;
            return this;
        }
        
        public Builder contextWindow(int contextWindow) {
            this.contextWindow = contextWindow;
            return this;
        }
        
        public Builder minSequenceLength(int minSequenceLength) {
            this.minSequenceLength = minSequenceLength;
            return this;
        }
        
        public Builder targetSequenceLength(int targetSequenceLength) {
            this.targetSequenceLength = targetSequenceLength;
            return this;
        }
        
        public SamplingConfig build() {
            return new SamplingConfig(
                temperature,
                topP,
                topK,
                repetitionPenalty,
                doSample,
                eosTokenIDs,
                contextWindow,
                minSequenceLength,
                targetSequenceLength
            );
        }
    }
    
    // Getters
    public float getTemperature() {
        return temperature;
    }
    
    public float getTopP() {
        return topP;
    }
    
    public int getTopK() {
        return topK;
    }
    
    public float getRepetitionPenalty() {
        return repetitionPenalty;
    }
    
    public boolean isDoSample() {
        return doSample;
    }
    
    public long[] getEosTokenIDs() {
        return eosTokenIDs;
    }
    
    public int getContextWindow() {
        return contextWindow;
    }
    
    public int getMinSequenceLength() {
        return minSequenceLength;
    }
    
    public int getTargetSequenceLength() {
        return targetSequenceLength;
    }
}

